import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import logging
import os

# Configure logging
logging.basicConfig(level=logging.WARNING, format='%(asctime)s - %(levelname)s - %(message)s')

def test_all_stocks(folder):
    hold_5_days = []
    hold_10_days = []
    hold_20_days = []
    hold_30_days = []
    hold_60_days = []
    for file in os.listdir(folder):
        if file.endswith('.csv'):
            path = os.path.join(folder, file)
            result = test_stock(path)
            hold_5_days.append(result[0])
            hold_10_days.append(result[1])
            hold_20_days.append(result[2])
            hold_30_days.append(result[3])
            hold_60_days.append(result[4])

    print(f"Average return for holding 5 days: {np.mean(hold_5_days):.4f}")
    print(f"Average return for holding 10 days: {np.mean(hold_10_days):.4f}")
    print(f"Average return for holding 20 days: {np.mean(hold_20_days):.4f}")
    print(f"Average return for holding 30 days: {np.mean(hold_30_days):.4f}")
    print(f"Average return for holding 60 days: {np.mean(hold_60_days):.4f}")



def test_stock(path):
    """
    Parameters:
    path (str): The file path to the stock data CSV file.

    Returns:
    None
    """
    # 获取股票数据
    data = pd.read_csv(path, parse_dates=['date'])

    # 获取股票数据
    data = pd.read_csv(path)

    # 确保索引是日期时间类型
    if 'date' in data.columns:
        data.index = pd.to_datetime(data.date)
    else:
        raise ValueError("The 'date' column is missing from the data.")

    # 确保索引是日期时间类型
    data.index = pd.to_datetime(data.date)

    # 计算5日均线
    data['5_day_mavg'] = data['close'].rolling(window=5).mean()

    # 计算10日均线
    data['10_day_mavg'] = data['close'].rolling(window=10).mean()

    # 判断下跌趋势
    data['20_day_mavg'] = data['close'].rolling(window=20).mean()
    data['falling_trend'] = (
        (data['close'] < data['20_day_mavg']) &
        (data['close'].shift(1) < data['20_day_mavg'].shift(1)) &
        (data['close'].shift(2) < data['20_day_mavg'].shift(2)) &
        (data['close'].shift(3) < data['20_day_mavg'].shift(3)) &
        (data['close'].shift(4) < data['20_day_mavg'].shift(4)) &
        (data['close'].shift(5) < data['20_day_mavg'].shift(5)) &
        (data['close'].shift(6) < data['20_day_mavg'].shift(6)) &
        (data['close'].shift(7) < data['20_day_mavg'].shift(7)) &
        (data['close'].shift(8) < data['20_day_mavg'].shift(8)) &
        (data['close'].shift(9) < data['20_day_mavg'].shift(9)) &
        (data['close'].shift(10) < data['20_day_mavg'].shift(10)) &
        (data['5_day_mavg'] < data['10_day_mavg']) & 
        (data['5_day_mavg'].shift(1) < data['10_day_mavg'].shift(1)) & 
        (data['5_day_mavg'].shift(2) < data['10_day_mavg'].shift(2)) & 
        (data['5_day_mavg'].shift(3) < data['10_day_mavg'].shift(3)) & 
        (data['5_day_mavg'].shift(4) < data['10_day_mavg'].shift(4)) & 
        (data['5_day_mavg'].shift(5) < data['10_day_mavg'].shift(5)) & 
        (data['5_day_mavg'].shift(6) < data['10_day_mavg'].shift(6)) & 
        (data['5_day_mavg'].shift(7) < data['10_day_mavg'].shift(7)) & 
        (data['5_day_mavg'].shift(8) < data['10_day_mavg'].shift(8)) &
        (data['5_day_mavg'].shift(9) < data['10_day_mavg'].shift(9)) &
        (data['5_day_mavg'].shift(10) < data['10_day_mavg'].shift(10)) &
        (data['10_day_mavg'] < data['20_day_mavg']) & 
        (data['10_day_mavg'].shift(1) < data['20_day_mavg'].shift(1)) & 
        (data['10_day_mavg'].shift(2) < data['20_day_mavg'].shift(2)) & 
        (data['10_day_mavg'].shift(3) < data['20_day_mavg'].shift(3)) & 
        (data['10_day_mavg'].shift(4) < data['20_day_mavg'].shift(4)) & 
        (data['10_day_mavg'].shift(5) < data['20_day_mavg'].shift(5)) & 
        (data['10_day_mavg'].shift(6) < data['20_day_mavg'].shift(6)) & 
        (data['10_day_mavg'].shift(7) < data['20_day_mavg'].shift(7)) & 
        (data['10_day_mavg'].shift(8) < data['20_day_mavg'].shift(8)) &
        (data['10_day_mavg'].shift(9) < data['20_day_mavg'].shift(9)) &
        (data['10_day_mavg'].shift(10) < data['20_day_mavg'].shift(10))

        )

    # 寻找放量阳线
    data['20_day_avg_volume'] = data['volume'].rolling(window=20).mean()
    data['bullish_volume'] = (data['volume'] > 2 * data['20_day_avg_volume']) & (data['close'] > data['10_day_mavg'])

    # 生成买入信号
    data['close_to_10_day_mavg'] = np.isclose(data['close'], data['10_day_mavg'], atol=0.03*data['10_day_mavg'])
    
    # 生成买入信号
    # data['buy_signal'] = data['falling_trend'].shift(8) & data['bullish_volume'].shift(8, fill_value=False) & data['close_to_10_day_mavg']

    data['buy_signal'] = (
        (data['falling_trend'].shift(5) & data['bullish_volume'].shift(4, fill_value=False) & data['close_to_10_day_mavg']) |
        (data['falling_trend'].shift(6) & data['bullish_volume'].shift(5, fill_value=False) & data['close_to_10_day_mavg']) |
        (data['falling_trend'].shift(7) & data['bullish_volume'].shift(6, fill_value=False) & data['close_to_10_day_mavg']) |
        (data['falling_trend'].shift(8) & data['bullish_volume'].shift(7, fill_value=False) & data['close_to_10_day_mavg']) |
        (data['falling_trend'].shift(9) & data['bullish_volume'].shift(8, fill_value=False) & data['close_to_10_day_mavg'])
    )

    # 检查生成的买入信号日期
    buy_signal_dates = data[data['buy_signal']].index
    # print(f"Buy signal dates: {buy_signal_dates}")
    dict = {}
    dict[os.path.basename(path).split('.')[0]] = [str(date)[:10] for date in buy_signal_dates]
    # print(dict)
    
    # 初始化用于记录不同持有期收益的DataFrame
    hold_periods = [5, 10, 20, 30, 60]
    hold_days = [f'hold_{period}_days' for period in hold_periods]
    returns_df = pd.DataFrame(index=data.index, columns=hold_days)

    # 计算不同持有期的收益
    for period in hold_periods:
        for buy_date in buy_signal_dates:
            sell_date_idx = data.index.get_loc(buy_date) + period
            if sell_date_idx < len(data.index):
                buy_price = data.loc[buy_date, 'close']
                sell_date = data.index[sell_date_idx]
                sell_price = data.loc[sell_date, 'close']
                if pd.isna(sell_price):
                    logging.warning(f"Sell price is NaN for buy date {buy_date} and sell date {sell_date}")
                    returns_df.loc[buy_date, f'hold_{period}_days'] = np.nan
                else:
                    returns_df.loc[buy_date, f'hold_{period}_days'] = (sell_price - buy_price) / buy_price
            else:
                # logging.warning(f"Sell date index {sell_date_idx} is out of range for buy date {buy_date}")
                returns_df.loc[buy_date, f'hold_{period}_days'] = np.nan

    # 绘制不同持有期的累计收益曲线
    results = []
    for period in hold_periods:
        cumulative_returns = (1 + returns_df[f'hold_{period}_days']).cumprod(skipna=True) - 1
        plt.plot(cumulative_returns, label=f'Hold {period} days')
        avg_return = returns_df[f'hold_{period}_days'].mean(skipna=True)
        results.append(avg_return)
        # print(f'Average return for holding {period} days: {avg_return:.4f}')
    
    # input_filename = os.path.basename(path)
    # output_filename = f'{os.path.splitext(input_filename)[0]}_processed.csv'
    # data.to_csv(output_filename)
    # input_filename = path.split('/')[-1].split('.')[0]
    # returns_df.to_csv(f'{input_filename}_returns_df.csv')

    # plt.title('Cumulative Returns for Different Holding Periods')
    # plt.xlabel('Date')
    # plt.ylabel('Returns')
    # plt.legend()
    # plt.show()
    return dict
    return results


def find_all_buys(path):
    all_buys = {}
    for file in os.listdir(path):
        if file.endswith('.csv'):
            dict = test_stock(os.path.join(path, file))
            for key, value in dict.items():
                for date in value:
                    if date in all_buys:
                        all_buys[date].append(key)
                    else:
                        all_buys[date] = [key]
        # break
    all_buys = sorted(all_buys.items())
    
    with open('all_buys.txt', 'w') as f:
        for date, stocks in all_buys:
            f.write(f'{date}: {", ".join(stocks)}\n')
        # print(f'{date}: {stocks}')


def save_dict_to_csv(dict, filename):
    df = pd.DataFrame(dict)
    df.to_csv(filename, index=False)


if __name__ == '__main__':
    # print(test_stock('stocks/sz300763.csv'))
    # print(test_stock('stocks/sz301153.csv'))
    # print(test_stock('stocks/sz301439.csv'))
    # find_all_buys('stocks_part')
    find_all_buys('stocks')
